"""Network :: distributed estimation"""

import numpy as np

class NetworkNode(object):
    """
    Node of a network

    Attributes
    ----------
    model : pybamo.regmodel or child
        Attached model
    neighbours : list
        List of adjacent neighbours
        
        Each neighbour is a dictionary:
        neighbour = {'no': int, 'incremental_weight': float, 'spatial_weight': float}
        
        where 'no' is a number of node, 'incremental_weight' and 'spatial_weight' are weights for updates
    current_data : numpy.array
        Last assigned data vector


    Methods
    -------
    __init__(model, neighbours, desc='Node')
        Constructor 
    set_own_data(data)
        Assign 'data' to 'current_data'
    incremental_update()
        Run incremental update
    spatial_update
        Run spatial update
    """

    def __init__(self, model, desc='Node'):
        """
        Constructor

        Parameters
        ----------
        model : pybamo.regmodel or child
            Model
        desc : str
            Description
        netighbours : dict
            Neighbours
        """
        self._model = model
        self._neighbours = []
        self.desc = desc
        self._data = 0

    def set_neighbours(self, neighbours):
        self._neighbours = neighbours

    def set_own_data(self, data):
        """
        Assign data to self.data

        Parameters
        ----------
        data : numpy.array
        """
        self._data = data

    def incremental_update(self):
        """
        Run incremental update
        """
        for neighbour in self._neighbours:
            data = neighbour['node'].data
            weight = neighbour['incremental_weight']
            self._model.model_update(data * np.sqrt(weight))

    def spatial_update(self):
        """
        Run spatial update
        """
        global_est_theta = 0
        for neighbour in self._neighbours:
            est_theta = neighbour['node']._model.est_theta
            weight = neighbour['spatial_weight']
            global_est_theta += weight * est_theta
        return global_est_theta

    @property
    def data(self):
        """Returns data"""
        return self._data



class Network(object):
    """
    Network with distributed estimation

    Attributes
    ----------
    nodes : dictionary
        Dictionary of nodes {identificator : Network.NetworkNode}
    _log : dictionary
        Log in the form {identificator : [arrays of est. coeffs.]}
    log_list_by_nodes : list
        Log, np.arrays in a list
    log_3d_array : numpy.array
        Log, 3D np.array
        rows : time instants `t'
        columns : estimates (vectors)
        3rd axis : nodes

    Methods
    -------
    __init__()
        Constructor
    init_log()
        Create self._log
    update(network_data)
        Update NetworkNodes._data and run incremental update
    estimate()
        Runs spatial update
    stack_results()
        Fills self.log_list_by_nodes and self.log_3d_array
    """

    def __init__(self):
        self.nodes = {}     #{'1': node1, '2': node2,...}
        self._log = {}
        self.log_list_by_nodes = []
        self.log_3d_array = 0

    def init_log(self):
        for ident in self.nodes.iterkeys():
            self._log[ident] = []

    def update(self, network_data):
        #network_data = {'1': data1, '2': data2,...}
        if not self._log:
        	self.init_log()

        # Put data to nodes
        for ident, data in network_data.iteritems():
            self.nodes[ident].set_own_data(data)

        # Run incremental update
        for ident in self.nodes.iterkeys():
        	self.nodes[ident].incremental_update()

    def estimate(self):
        #spatial_update()
        for ident in self.nodes.iterkeys():
        	est = self.nodes[ident].spatial_update()
        	self._log[ident].append(est)

    def stack_results(self):
        for ident in self._log.iterkeys():
            self.log_list_by_nodes.append(np.vstack(self._log[ident]))
        self.log_3d_array = np.array(self.log_list_by_nodes)


#--------------------------------------------

if __name__ == "__main__":
    from copy import deepcopy
    from regmodel import *

    model1 = Regmodel(np.eye(3), 5)
    node1 = NetworkNode(model1, desc='1st node')

    model2 = deepcopy(model1)
    node2 = NetworkNode(model2, desc='2nd node')

    neigh_of_node1 = []
    neigh_of_node1.append({'node': node1, 'incremental_weight': 1., 'spatial_weight': 1.})
    neigh_of_node1.append({'node': node2, 'incremental_weight': 0., 'spatial_weight': 0.})
    node1.set_neighbours(neigh_of_node1)

    neigh_of_node2 = []
    neigh_of_node2.append({'node': node1, 'incremental_weight': 0.3, 'spatial_weight': 0.5})
    neigh_of_node2.append({'node': node2, 'incremental_weight': 0.7, 'spatial_weight': 0.5})
    node2.set_neighbours(neigh_of_node2)

    #data1 = np.array([1, 2, 1])
    #node1.set_own_data(data1)

    #data2 = np.array([1, 2, 1])
    #node2.set_own_data(data2)

    #print 'Before incremental update (node1)'
    #print node1._model.v.v
    #node1.incremental_update()
    #print 'After incremental update (node1)'
    #print node1._model.v.v
    #print ''
    #print 'Before incremental update (node2)'
    #print node2._model.v.v
    #node2.incremental_update()
    #print 'After incremental update (node2)'
    #print node2._model.v.v

    #print 'Global estimated parameter (node1)'
    #print node1.spatial_update()
    #print ''
    #print 'Global estimated parameter (node2)'
    #print node2.spatial_update()

    network = Network()
    network.nodes = {
    	    'node1': node1,
            'node2': node2
    }

    data = {
    	    'node1': np.array([1,2,1]),
    	    'node2': np.array([1,2,0])
    }

    for i in xrange(10):
        network.update(data)
        network.estimate()

    network.stack_results()
